#include "image_op.h"
#include <iostream>

namespace vkinfer
{
	AffineMatrix::AffineMatrix(const Rect& src, const Rect& dst)
	{
		m[0][1] = 0.0f;
		m[1][0] = 0.0f;
		float scale_x = dst.width * 1.0f / src.width;
		float scale_y = dst.height * 1.0f / src.height;
		m[0][0] = scale_x;
		m[1][1] = scale_y;
		m[0][2] = dst.x - scale_x * src.x;
		m[1][2] = dst.y - scale_y * src.y;
	}

	AffineMatrix::AffineMatrix(uint32_t src_width, uint32_t src_height, uint32_t dst_width, uint32_t dst_height)
	{
		m[0][1] = 0.0f;
		m[1][0] = 0.0f;
		float scale_x = dst_width * 1.0f / src_width;
		float scale_y = dst_height * 1.0f / src_height;
		m[0][0] = scale_x;
		m[1][1] = scale_y;
		m[0][2] = 0.0f;
		m[1][2] = 0.0f;
	}

	AffineMatrix AffineMatrix::inverse() const
	{
		AffineMatrix result;
		float det = m[0][0] * m[1][1] - m[0][1] * m[1][0];

		result.m[0][0] = m[1][1] / det;
		result.m[0][1] = -m[0][1] / det;
		result.m[0][2] = (m[0][1] * m[1][2] - m[1][1] * m[0][2]) / det;
		result.m[1][0] = -m[1][0] / det;
		result.m[1][1] = m[0][0] / det;
		result.m[1][2] = (m[1][0] * m[0][2] - m[0][0] * m[1][2]) / det;
		return result;
	}

	void AffineMatrix::apply(float x, float y, float* px, float* py)
	{
		*px = m[0][0] * x + m[0][1] * y + m[0][2];
		*py = m[1][0] * x + m[1][1] * y + m[1][2];
	}

	ImagePreprocessor::ImagePreprocessor()
	{
		width = 224;
		height = 224;
		channel_shuffle = false;
		hwc2chw = true;
		normalize = true;
		mean = { 0.5, 0.5, 0.5 };
		std = { 0.5, 0.5, 0.5 };
	}

	std::shared_ptr<Tensor<float>> ImagePreprocessor::process(const Image& image)
	{
		// cpu computing
		std::vector<uint32_t> shape;
		if (hwc2chw)
			shape = { 1, image.channels, height, width };
		else
			shape = { 1, height, width, image.channels };

		std::shared_ptr<Tensor<float>> tensor = std::make_shared<Tensor<float>>(shape);

		Rect src_rect = { 0, 0, width, height };
		Rect dst_rect = { 0, 0, image.width, image.height };
		AffineMatrix affine(src_rect, dst_rect);
		for (uint32_t y = 0; y < height; ++y)
		{
			for (uint32_t x = 0; x < width; ++x)
			{
				float ox, oy;
				affine.apply((float)x, (float)y, &ox, &oy);
				for (uint32_t c = 0; c < image.channels; ++c)
				{
					// bilinear interpolate
					float v = image.sample(ox, oy, c);
					if (normalize)
						v = v / 255.0f;

					uint32_t oc = c;
					if (channel_shuffle)
						oc = 2 - c;
					uint32_t dst_index;
					if (hwc2chw)
						dst_index = oc * width * height + y * width + x;
					else
						dst_index = y * width * image.channels + x * image.channels + oc;
					
					tensor->data[dst_index] = (v - mean[oc]) / std[oc];
				}
			}
		}

		return tensor;
	}

	void display_image(const Image& image)
	{
		std::cout << "no display context available\n";
	}
}